#include "CMAOptimizer.h"
#include <cmath>
#include <memory.h>
#include <algorithm>

using namespace std;
using namespace learn;

learn::CMAOprimizer::CMAOprimizer()
	: isFirstIteration(true) 
{	
}

learn::CMAOprimizer::~CMAOprimizer() {
}

template <typename T>
inline T clamp(T x, T minVal, T maxVal) {
	return min(max(minVal, x), maxVal);
}

void copyVector(float *dest, float *src, size_t count) {
	memcpy(dest, src, count * sizeof(float));
}

void learn::CMAOprimizer::init(Model &model, CMAMode cmaMode, int lambda, int mu) {
	
	this->cmaMode = cmaMode;
	
	// TODO rename	
	int dim = model.getParameterCount();
	if (cmaMode == RANK_MU_UPDATE || cmaMode == RANK_ONE_UPDATE) {
		if (lambda <= 0) {
			lambda = CMA::suggestLambda(dim);
		}
		if (mu <= 0) {
			mu = CMA::suggestMu(lambda, CMA::RecombType::superlinear);
		}
				
		vector<float> chromW(dim);
		model.getParameters(&chromW[0], 0, dim);
				
		// fill parents population
		parentParams.resize(mu * dim);
		parents.resize(mu);				
		for (int i = 0; i < mu; i++) {
			parents[i].parameters = &parentParams[i * mu];
			copyVector(parents[i].parameters, &chromW[0], dim);
		}
		
		// fill offspring population
		offspringParams.resize(lambda * dim);
		offsprings.resize(lambda);		
		for (int i = 0; i < lambda; i++) {
			offsprings[i].parameters = &offspringParams[i * dim];
			copyVector(offsprings[i].parameters, &chromW[0], dim);
		}
		
		if (cmaMode == RANK_MU_UPDATE) {
			cma.init(dim,  0.01f, 1.0f, 
					&parents[0], parents.size(),
					CMA::superlinear, CMA::rankmu);		
		} else {
			cma.init(dim,  0.01f, 1.0f, 
					&parents[0], parents.size(),
					CMA::superlinear, CMA::rankone);
		}
	} 
	else if (cmaMode == ONE_PLUS_ONE) { 
		// TODO
	}
		
	isFirstIteration = true;	
}
	
float getError(Model& model, ErrorFunction &errorFunction, OptimizationSupervisor &data) {
	// TODO set parameters
	// TODO get error
	return 0;
}

float learn::CMAOprimizer::optimize(Model &model, ErrorFunction &errorFunction, OptimizationSupervisor &data) {
	// TODO rename
	int dim = model.getParameterCount();
		
	if (isFirstIteration) {
		// ensure that the parents have a valid fitness value
		for (int i = 0; i < parents.size(); i++) {
			model.setParameters(parents[i].parameters, 0, dim);
			parents[i].fitness = errorFunction.getError(model, data);
		}
		isFirstIteration = false;
	}
		
	if (cmaMode == RANK_MU_UPDATE || cmaMode == RANK_ONE_UPDATE) {		
		// TODO parallel		
		for (size_t i = 0; i < offsprings.size(); i++) {
			cma.createIndividual(offsprings[i].parameters);
			
			model.setParameters(offsprings[i].parameters, 0, dim);
			offsprings[i].fitness = errorFunction.getError(model, data);
		}
		
		// TODO implement uncertaintyHandling
		
		// pick parents.size() best individuals
		sort(offsprings.begin(), offsprings.end());
		parents.assign(offsprings.begin(), offsprings.begin() + parents.size());
		
		cma.updateStrategyParameters(&parents[0], parents.size());
	}
	
	Individual bestIndivid = parents[0];
	if (bestIndivid.fitness < bestFitness) {
		copyVector(&bestParameters[0], bestIndivid.parameters, dim);
		bestFitness = bestIndivid.fitness;
	}
	
	model.setParameters(&bestParameters[0], 0, dim);
	return bestFitness;
}

